# -*- coding: utf-8 -*-  
'''
GSDMM相关零件

Created on 2021年9月2日
@author: luoyi
'''
import tensorflow as tf

import utils.conf as conf
from utils.iexicon import LiteWordsWarehouse
from models.gsdmm.process import log_topic_distribution, incr_num_topic_doc, stat_nw_docs


#    Gibbs层
class GibbsLayer(tf.keras.layers.Layer):
    '''Gibbs采样层，原始论文中的方法
        初始化：
            - 主题数：K
            - 词库总词数：V
            - 两个Dirichlet分布超参：α={50./K}，β={0.2}
            - 主题-词 计数：wn_topic_word = Tensor(K, V)，初始值为1（防止除零错误）
            - 主题-文档 计数：dn_topic_doc = Tensor(K)，初始值为1
        前向过程：
            按照边缘概率计算每篇文档的所属每个主题概率：
                原论文的算法，但连乘时文档长一点就会区域无穷
                P(z[i]=k | W,α,β) = dn_topic_doc[k] + α[k] / ∑(K)( dn_topic_doc[k] + α[k] )
                                        *
                                    ∏(w ∈ W[d])( wn_topic_word[k,w] + β[w] ) /  ∏(j=1->n[d])( ∑(V)(wn_topic_word[k,v] + β[v]) + j - 1 )
                logP(z[i]=k | W,α,β) = log( dn_topic_doc[k] + α[k] )
                                        - log( ∑(K)( dn_topic_doc[k] + α[k] ) )
                                        + ∏(w ∈ W[d])log( wn_topic_word[k,w] + β[w] )                    ------ W[d]表示第d篇文档中的词集合，可能会有重复的词
                                        - ∑(j=1->n[d])log( ∑(V)(wn_topic_word[k,v] + β[v]) + j - 1 )
                计算文档从属每个主题的概率，组成多项分布θ
            多项分布θ中采样拿到文档的new_topic
            根据new_topic更新相关计数
            wn_topic_word[new_topic, v] += n[d, v]
            dn_topic_doc[new_topic] += 1
    '''
    def __init__(self,
                 name='gibbs_layer',
                 K=conf.GSDMM.get_k(),
                 V=LiteWordsWarehouse.instance().words_count(), 
                 **kwargs
                 ):
        self._K = K
        self._V = V
        
        super(GibbsLayer, self).__init__(name=name, **kwargs)
        pass
    
    def build(self, input_shape):
        #    初始化α超参 Tensor(K,)
        self._alpha = self.add_weight(name='gibbs/dirichlet/alpha', 
                                      shape=(self._K, ), 
                                      dtype=tf.float32, 
#                                       initializer=lambda shape, dtype: tf.convert_to_tensor([50. / self._K] * self._K),
                                      trainable=False)
        self._alpha.assign(tf.convert_to_tensor([50. / self._K] * self._K))
        #    初始化超参β Tensor(V,)
        self._beta = self.add_weight(name='gibbs/dirichlet/beta', 
                                     shape=(self._V, ), 
                                     dtype=tf.float32, 
                                     trainable=False)
        self._beta.assign(tf.convert_to_tensor([0.2] * self._V))
        
        #    初始化主题-词 计数 Tensor(K, V)
        self._wn_topic_word = self.add_weight(name='gibbs/wn_topic_word', 
                                              shape=(self._K, self._V), 
                                              dtype=tf.float32, 
                                              initializer=tf.initializers.Zeros(), 
                                              trainable=False)
#         self._wn_topic_word.assign(tf.ones(shape=(self._K, self._V)) * 0.00001)
        #    初始化主题-文档 计数 Tensor(K)
        self._dn_topic_doc = self.add_weight(name='gibbs/dn_topic_doc', 
                                             shape=(self._K, ), 
                                             dtype=tf.float32, 
                                             initializer=tf.initializers.Zeros(), 
                                             trainable=False)
#         self._dn_topic_doc.assign(tf.ones(shape=(self._K), ) * 0.0001)
        
        #    初始化主题-词 分布    Tensor(K, V)
        self._prob_topic_word = self.add_weight(name='gibbs/prob_topic_word', 
                                                shape=(self._K, self._V), 
                                                dtype=tf.float32, 
                                                initializer=tf.initializers.zeros(), 
                                                trainable=False)
        #    初始化主题-文档 分布 Tensor(K)
        self._prob_doc_topic = self.add_weight(name='gibbs/prob_doc_topic', 
                                               shape=(self._K, ), 
                                               dtype=tf.float32, 
                                               initializer=tf.initializers.zeros(), 
                                               trainable=False)
        
        super(GibbsLayer, self).build(input_shape)
        pass
    
    #    计算当前的主题-词分布：φ
    def op_topic_word_prob(self):
        p = self._wn_topic_word / tf.math.reduce_sum(self._wn_topic_word, axis=-1)[:, tf.newaxis]
        self._prob_topic_word.assign(p)
        pass
    #    取当前 主题-词 分布：φ[k,v]    Tensor(K, V)
    def topic_word_prob(self):
        return self._prob_topic_word
    #    取当前 主题-词 计数    Tensor(K, V)
    def wn_topic_word(self):
        return self._wn_topic_word
    #    计算当前的主题-文档分布：θ
    def op_doc_topic_prob(self):
        p = self._dn_topic_doc / tf.math.reduce_sum(self._dn_topic_doc)
        self._prob_doc_topic.assign(p)
        pass
    #    取当前 主题-文档 分布：θ    Tensor(K,)
    def docs_topic_prob(self):
        return self._prob_doc_topic
    #    取当前 主题-文档 计数    Tensor(K,)
    def wn_docs_topic(self):
        return self._dn_topic_doc
    
    
    def call(self, inputs, training=None):
        '''根据边缘概率计算每篇文档的所属主题概率分布
            logP(z[i]=k | W,α,β) = log( dn_topic_doc[k] + α[k] )
                                    - log( ∑(K)( dn_topic_doc[k] + α[k] ) )
                                    + ∏(w ∈ W[d])log( wn_topic_word[k,w] + β[w] )                    ------ W[d]表示第d篇文档中的词集合，可能会有重复的词
                                    - ∑(j=1->n[d])log( ∑(V)(wn_topic_word[k,v] + β[v]) + j - 1 )
                                
            @param inputs: Tensor(batch_size, max_doc_words)    -1表示填充词
        '''
        #    计算每篇文档的主题概率    Tensor(batch_size, K)
        p = log_topic_distribution(K=self._K,
                                   docs=inputs,
                                   dn_topic_doc=self._dn_topic_doc,
                                   wn_topic_word=self._wn_topic_word,
                                   alpha=self._alpha,
                                   beta=self._beta)
        #    采样每篇文档的新主题    Tensor(batch_size, )
        new_topics = tf.random.categorical(p, num_samples=1)
        new_topics = tf.squeeze(new_topics, axis=-1)
        self._crt_new_topics = new_topics
        
        #    更新计数
        self.update_count(new_topics, inputs)
        
        #    重新计算θ，φ
        if (training): 
            self.op_topic_word_prob()
            self.op_doc_topic_prob()
            pass
        
        return inputs
    

    
    #    更新计数
    def update_count(self, new_topics, docs):
        '''
            wn_topic_word[new_topic, v] += n[d, v]
            dn_topic_doc[new_topic] += 1
            
            @param new_topics: 每篇文档的新主题    Tensor(batch_size, )
            @param docs: 每篇文档的词id    Tensor(batch_size, max_docs_words)
        '''
        #    更新 主题-文档 计数
        self.update_num_topic_doc(new_topics)
        
        #    更新 主题-词 词计数
        self.update_num_topic_word(new_topics, docs)
        pass

    #    更新 主题-文档 计数
    def update_num_topic_doc(self, new_topics):
        #    计算new_topics中每个topic的数量
        t_incr = incr_num_topic_doc(new_topics, K=self._K)
        
        #    执行累加
        new_dn = self._dn_topic_doc + t_incr
        self._dn_topic_doc.assign(new_dn)
        pass
    #    更新 主题-词 的词计数
    def update_num_topic_word(self, new_topics, docs):
        '''
            @param new_topics: Tensor(batch_size, )
            @param docs: Tensor(batch_size, max_docs_words)
        '''
        #    统计每篇文档，在每个主题下，每个词的增量
        incrs = stat_nw_docs(new_topics, docs, K=self._K, V=self._V)                               #    Tensor(batch_size, K, V)
        incrs = tf.math.reduce_sum(incrs, axis=0)                                                  #    Tensor(K, V)
        incrs = tf.cast(incrs, dtype=tf.float32)
        
        #    累加增量
        self._wn_topic_word.assign(self._wn_topic_word + incrs)
        pass
    
    pass